'''
Created on Nov 9, 2012

@author: Anand
'''
from matplotlib import pyplot as plt
import numpy as na
import random


class Constants(object):
  def __init__(self):
    self.GEL_HEIGHT = 1000
    self.GEL_MARGIN = 30
    self.GEL_WIDTH = 500
    self.LANE_WIDTH = 80
    self.LANES = 5

CONS = Constants()

'''
 DNA on a gel separates by  y = c k ** -x 
 (y is the distance on gel from init well, and x is the oligonucleotide length)
'''

class GelMap(object):
  def __init__(self, gel_start, gel_end, dna_longest, dna_shortest):
    pass
  def range_embed(self, x):
    return None

class LinGelMap(GelMap):
  def __init__(self, gel_start, gel_end, dna_longest, dna_shortest):
    self.m = (gel_end-gel_start)/float(dna_shortest-dna_longest)
    self.b = gel_end - dna_shortest * ((gel_end-gel_start)/float(dna_shortest-dna_longest))
  def range_embed(self, x):
    return int(na.round(self.m*x + self.b))

class ExpGelMap(GelMap):
  def __init__(self, gel_start, gel_end, dna_longest, dna_shortest):
    self.c = gel_start *(( gel_start / float(gel_end) ) ** (dna_longest/float(dna_shortest-dna_longest)) )
    self.k = (gel_start/ float(gel_end) ) ** (1/float(dna_shortest - dna_longest))
  def range_embed(self, x):
    return int(na.round(self.c * (self.k)**(-x)))

class RangeStandard(object):
  def __init__(self):
    self.ladder = None
    self.map = None
  def get_position(self, x):
    return self.map.range_embed(x)

class LowRangeStandard(RangeStandard):
  def __init__(self):
    self.ladder = [10,15,20,25,35,50,75,100,150,200,300]
    self.map = ExpGelMap(CONS.GEL_MARGIN, CONS.GEL_HEIGHT-CONS.GEL_MARGIN, max(self.ladder), min(self.ladder))
  
class HighRangeStandard(RangeStandard):
  def __init__(self):
    self.ladder = [75,200,300,400,500,700,1000,1500,2000,3000,4000,5000,7000,10000]
    self.map = ExpGelMap(CONS.GEL_MARGIN, CONS.GEL_HEIGHT-CONS.GEL_MARGIN, max(self.ladder), min(self.ladder))


class TRAlleles(object):
  def __init__(self):
    self.offsets = []
    self.alleles = {}
    self.colors = []
  def default(self):
    self.alleles = {}
    self.colors = [1,1,1]
    self.offsets = [('D3S1358',20), ('D8S1179',10), ('FGA',20)]
    self.alleles['D3S1358'] = na.arange(15,18)*4
    self.alleles['D8S1179'] = na.arange(12,16)*4
    self.alleles['FGA'] = na.arange(22,26)*4

  def parse(self, fpath):
    """
    Each entry is <Loci>\t<Offset>\t<RepeatCounts>\t<RepeatLength>
    """
    with open(fpath, 'rb') as f:
      for l in f:
        tokens = l.strip().split('\t')
        loci = tokens[0]
        o,rc, rl = tokens[1:]
        self.offsets.append((loci,int(o)))
        self.alleles[loci] = na.array(map(int, rc.split(',')))*int(rl)
      self.colors = [1]*len(self.offsets)
  
class VNTRAlleles(TRAlleles):
  def __init__(self):
    TRAlleles.__init__(self)
  def default(self):
    self.alleles = {}
    self.colors = [1,1,-1]
    self.offsets = [('VNTR',300), ('D1S80',300), ('D17S79',1000)]
    self.alleles['VNTR'] = na.array([70,100,120,200,400,450])*17
    self.alleles['D1S80'] = na.array([24, 18, 29, 20, 33])*18
    self.alleles['D17S79'] = na.array([36,41,43,50])*38

  def parse(self, fpath):
    """
    Each entry is <Loci>\t<Offset>\t<BinaryColor(-1|1)>\t<RepeatCounts>\t<RepeatLength>
    """
    with open(fpath, 'rb') as f:
      for l in f:
        tokens = l.strip().split('\t')
        loci = tokens[0]
        o,c,rc, rl = tokens[1:]
        self.offsets.append((loci,int(o)))
        self.alleles[loci] = na.array(map(int, rc.split(',')))*int(rl)
        self.colors.append(int(c))

class GenerateTRMatch(object):
  def __init__(self, parent_flag=False, figfpath=None, band_width=5, standard=LowRangeStandard(), cm=plt.get_cmap('gray_r')): 
    self.parent_flag = parent_flag
    self.fig_fpath = figfpath
    self.BAND_HALFWIDTH = band_width
    self.standard = standard
    self.cm = cm
  def _add_title_text(self, ax):
    xpos = na.arange(0,CONS.GEL_WIDTH,CONS.GEL_WIDTH/CONS.LANES)+(CONS.GEL_WIDTH/(CONS.LANES*2))
    for p,n in zip(xpos,['CRIME\nSCENE', 'SUSPECT\nA','SUSPECT\nB','SUSPECT\nC', 'DNA\nladder']):
      ax.text(p,-10, n, horizontalalignment='center', verticalalignment='bottom', fontsize=10)
  
    ax.set_xticks([])
    ax.set_yticks([])
    
  def _add_range_ladder_text(self, ax):
    for p in self.standard.ladder:
      ax.text(CONS.GEL_WIDTH+5, self.standard.get_position(p), "%d"%p, horizontalalignment='left', verticalalignment='center', fontsize=10)
  
  def simulate(self, allele_obj):
    #random.seed(0)
    fig = plt.figure(figsize=(5,9))
    ax = fig.add_subplot(111)
    # generate 3 suspects
    lanes = []
    for i in xrange(3):
      lane = []
      for k,off in allele_obj.offsets:
        lane.append(random.choice(allele_obj.alleles[k])+off)
        lane.append(random.choice(allele_obj.alleles[k])+off)  
      lanes.append(lane)
    
    # lane marker colors
    lane_colors = ['']*(2*len(allele_obj.offsets))
    lane_colors[::2] = allele_obj.colors
    lane_colors[1::2] = allele_obj.colors
      
    # choose a suspect or parent
    if self.parent_flag:
      i,j = random.sample(range(3),2)
      crime_scene = ['']*len(lanes[i])
      crime_scene[::2] = lanes[i][::2]
      crime_scene[1::2] = lanes[j][::2]
      
      lanes = [crime_scene] + lanes + [self.standard.ladder]
      
    else:
      i = random.randint(0,2)
      lanes = [lanes[i]] + lanes + [self.standard.ladder]
  
    lcs = [lane_colors]*(len(lanes)-1) + [[1]*len(self.standard.ladder)]
    
    # Adds speckles 
    gel = na.random.rand(CONS.GEL_HEIGHT,CONS.GEL_WIDTH)/10.
    
    xpositions = na.arange(0,CONS.GEL_WIDTH,CONS.GEL_WIDTH/CONS.LANES)+10
    for xpos, l, lc in zip(xpositions, lanes, lcs):
      for y,c in zip(l,lc):
        mid_y = int(self.standard.get_position(y))
        min_band, max_band = max(min(mid_y-self.BAND_HALFWIDTH,CONS.GEL_HEIGHT),0), min(max(mid_y+self.BAND_HALFWIDTH,0),CONS.GEL_HEIGHT) 
        try:
          gel[min_band:max_band,xpos:xpos+CONS.LANE_WIDTH] = c*(na.random.rand(max_band-min_band,CONS.LANE_WIDTH)+0.5)
        except:
          pass
          
    ax.imshow(gel,cmap=self.cm, origin='upper', interpolation='gaussian', vmin=-1.5, vmax=1.5, aspect='equal')
    
    self._add_title_text(ax)
    self._add_range_ladder_text(ax)
    if self.fig_fpath is None:
      plt.show()
    else:
      plt.savefig(self.fig_fpath)

def generate_str_match():
  x = GenerateTRMatch(parent_flag=True)
  y = TRAlleles()
  y.default()
  x.simulate(y)

def generate_vntr_match():
  x = GenerateTRMatch(parent_flag=False, figfpath=None, band_width=3, standard=HighRangeStandard(), cm=plt.get_cmap('PuOr_r'))
  y = VNTRAlleles()
  y.default()
  x.simulate(y) 

if __name__ == '__main__':
  import sys
  import argparse
  
  parser = argparse.ArgumentParser(description='Generate crime scene DNA evidence.')

  parser.add_argument('-d, --debug', dest='debug', action='store_const', const=True, default=None, help='Displays a default STR and VNTR allele set')
  parser.add_argument('--vntr', dest='vntr', action='store_const', const=True, default=False, help='Creates a Southern Blot for VNTR alleles (requires VNTR text file).')
  parser.add_argument('--str', dest='tr', action='store_const', const=True, default=False, help='Creates a Electrophoregram for STR alleles (requires STR text file).')
  parser.add_argument('-p, --parent-flag', dest='parent_flag', action='store_const', const=True, default=False, help='Crime Scene sample is the progeny of two suspects.')
  parser.add_argument('alleles', type=str, nargs='?', default=None, help='TAB-delimited alleles text file. A VNTR row is of the form:<Loci> <Offset> <BinaryColor(-1|1)> <RepeatCounts(1,2,3)> <RepeatLength>\n\n A STR allele row is of the form:<Loci> <Offset> <RepeatCounts(1,2,3)> <RepeatLength>')
  parser.add_argument('figfpath', type=str, nargs='?', default=None, help='Outputs figure to a file.')
  
  args = parser.parse_args()
  
  if args.debug is not None:
    generate_str_match()
    generate_vntr_match()
  elif args.alleles is not None:
    if args.vntr:
      x = GenerateTRMatch(parent_flag=args.parent_flag, figfpath=args.figfpath, band_width=3, standard=HighRangeStandard(), cm=plt.get_cmap('PuOr_r'))
      y = VNTRAlleles()
    elif args.tr:
      x = GenerateTRMatch(parent_flag=args.parent_flag, figfpath=args.figfpath, band_width=5, standard=LowRangeStandard(), cm=plt.get_cmap('gray_r'))
      y = TRAlleles()
    else:
      print >>sys.stderr, "Nothing to do, choose a --str or --vntr options"    
      sys.exit(1)
  
    y.parse(args.alleles)
    x.simulate(y)
  else:
    print >>sys.stderr, "Nothing to do, no allele text file supplied"
  
  
  #all_out()
  pass